import os

import cv2
import pytesseract
from sklearn.cluster import DBSCAN

from ocr_modules.pytesseract.tools import get_ocr_result_path, organize_line
from utils import utils, ocr_utils


# 整理所有行的数据（传入模型结果集，和用来定位的列表）
def organize_all_line_v1(char_boxe_list: list, original_image):
    # 获取竖线（用来给单元格定位用）
    col_list = ocr_utils.get_column_line(original_image.copy())

    # 存储所有行的数据 [ [line_cell_list] ,[line_cell_list] ,...]
    lines_cell_list = []
    # 存放每行的单元格 [ cell ,cell ,...]，由竖线的条数来判断
    line_cell_list = [[] for _ in range(len(col_list) - 1)]

    # 从第0条竖线开始
    col_index = 0

    for i in range(len(char_boxe_list)):
        p = col_list[col_index]  # 指针p用来标记起始位置
        box = char_boxe_list[i]
        c, x1, y1, x2, y2 = box[0], box[1], box[2], box[3], box[4]
        q = x1  # q指针用来标记当前位置

        # 如果当前位置在p指针的后面
        if q > p:
            # 查看在哪一列，遍历竖线，注意这里从 0 到 len-1
            for j in range(col_index, len(col_list) - 1):
                # 在第j条竖线和第j+1条竖线之间
                if col_list[j] < q < col_list[j + 1]:
                    # 移动指针
                    col_index = j
                    p = col_list[col_index]
                    break
            # 将字符对应的方框加到对应的单元格
            line_cell_list[col_index].append(box)

        # 如果在前面 到达了下一列
        else:
            # 整理上一行
            new_line_cell_list = organize_line(line_cell_list)
            # 将上一行加入到列表当中
            lines_cell_list.append(new_line_cell_list)
            # 重置数据
            line_cell_list = [[] for _ in range(len(col_list) - 1)]
            # 重新从下一行开始
            col_index = 0

        # 移动指针
        p = col_list[col_index]
    # 添加最后一行的数据
    new_line_cell_list = organize_line(line_cell_list)
    lines_cell_list.append(new_line_cell_list)
    return lines_cell_list


# 提取表格数据的函数（方法1）
def extract_table_data_v1(filepath: str):
    # 获取原始 OCR 结果文件路径
    ocr_path = get_ocr_result_path(filepath, "csv", "v1")
    ojr_path = get_ocr_result_path(filepath, "json", "v1")

    # 如果原始 OCR 结果文件已存在，则跳过扫描并打印相应信息
    if os.path.exists(ocr_path):
        print("{} 跳过扫描，有之前生成的扫描文本".format(filepath))
        return

    # 读取一张黑白图片
    original_image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)

    # OCR识别
    char_boxes = [box.split(" ") for box in
                  pytesseract.image_to_boxes(original_image, config=r"--oem 3 --psm 6 -l digits+eng").splitlines()]
    # 获取单个字符框的盒子
    char_boxs_list = [[box[0], int(box[1]), int(box[2]), int(box[3]), int(box[4])]
                      for box in char_boxes if box[0].isdigit() or box[0] == "."]

    # 将数据进行加工处理
    lines_cell_list = organize_all_line_v1(char_boxs_list, original_image)

    # # 下面这段代码可以在图像中框选数据
    # img = cv2.cvtColor(original_image, cv2.COLOR_GRAY2BGR)
    # height, width = img.shape[:2:]
    # for i in lines_cell_list:
    #     for j in i:
    #         x1, y1 = j[1]
    #         x2, y2 = j[3]
    #         # 在图像上绘制矩形框，
    #         x, y, w, h = j[0], j[1], j[2], j[3]
    #         cv2.rectangle(img, (x1, height - y1), (x2, height - y2), (255, 0, 255), 5)
    #         ocr_utils.cv_show("img", img)

    # 将结果转换为 JSON
    result_json = utils.list_to_json(lines_cell_list)

    # 保存文件
    utils.save_csv_file(data=result_json, ocr_path=ocr_path, ojr_path=ojr_path)

    print("{} pytesseract 扫描完成".format(ocr_path))


# 第二种方法的思想是画线法
def organize_all_line_v2(char_boxs_list: list, original_image):
    # 获取竖线（用来给单元格定位用）
    col_list = ocr_utils.get_column_line(original_image.copy())

    # 存放中点坐标
    char_centers = [[(box[1] + box[3]) // 2, (box[2] + box[4]) // 2] for box in char_boxs_list]

    # 聚合
    labels = DBSCAN(eps=50, min_samples=1).fit(char_centers).labels_

    # 下面的代码用来查看聚合函数的效果
    # 生成随机颜色
    import random
    rand_colors = []
    for i in range(len(labels)):
        r = random.randint(0, 255)
        g = random.randint(0, 255)
        b = random.randint(0, 255)
        rand_colors.append((r, g, b))
    img = cv2.cvtColor(original_image, cv2.COLOR_GRAY2BGR)
    height, width = img.shape[:2:]
    # 显示划线和方框
    for i in range(len(labels)):
        box = char_boxs_list[i]
        color = rand_colors[labels[i]]
        c, x1, y1, x2, y2 = box[0], box[1], box[2], box[3], box[4]
        cv2.rectangle(img, (x1, height - y1), (x2, height - y2), color, 5)
    ocr_utils.cv_show("img", img)

    # 存放所有单元格的列表
    cell_list = [[] for _ in range(max(labels) + 1)]
    for i in range(len(char_boxs_list)):
        cell_list[labels[i]].append(char_boxs_list[i])

    # 将单个字符组合成一个框
    lines_cell_list = []
    for i in cell_list:
        lines_cell_list.append(utils.group_cell(i))

    # 画线获得规定的 json 格式列表
    utils.export_all_line_data(original_image, lines_cell_list, col_list)


# 执行 OCR 过程的函数（方法2）
def extract_table_data_v2(filepath: str):
    # 获取原始 OCR 结果文件路径
    ocr_path = get_ocr_result_path(filepath, file_type="csv", append="v2")
    ojr_path = get_ocr_result_path(filepath, file_type="json", append="v2")

    # 如果原始 OCR 结果文件已存在，则跳过扫描并打印相应信息
    if os.path.exists(ocr_path):
        print("{} 跳过扫描，有之前生成的扫描文本".format(filepath))
        return

    # 读取一张黑白图片
    original_image = cv2.imread(filepath, cv2.IMREAD_GRAYSCALE)

    # OCR识别
    char_boxes = [box.split(" ") for box in
                  pytesseract.image_to_boxes(original_image, config=r"--oem 3 --psm 6 -l digits+eng")
                  .splitlines()]

    # 获取单个字符框的盒子
    char_boxs_list = [[box[0], int(box[1]), int(box[2]), int(box[3]), int(box[4])]
                      for box in char_boxes
                      if box[0].isdigit() or box[0] == "."]

    # 将数据进行加工处理
    lines_cell_list = organize_all_line_v2(char_boxs_list, original_image)

    # 将结果转换为 JSON
    result_json = utils.list_to_json(lines_cell_list)

    # 保存csv文件
    utils.save_csv_file(data=result_json, ocr_path=ocr_path, ojr_path=ojr_path)

    print("{} pytesseract 扫描完成".format(ocr_path))
